package net.pragyah.scalby.inference.exact

import net.pragyah.scalby.Variable

import net.pragyah.scalgorithms.graphs._

/*
 *  Pearl's Message passing algorithm 
 */
class MessagePassing_Tree[A](val G:Graph[Variable[A]],P : ProbabilityDistribution[A], val root:Vertex[Variable[A]]) {
  
  var aA = Set[Variable[A]]()
  
  assume(isTree)
  
  def isTree() : boolean = {
    //TODO .... verify if this is even a tree or not .... just do a depth first traversal and mark the visited nodes
    
    
    
    true
  }
  
  def initial_tree(){
    
    
    

	  
	  G.vertices.foreach(X =>{
	    X.tag = VertexTag[A](X.data)
	    X.getEdges.foreach(e =>{
	      val Y = e.getOther(X)
	      e.tag = EdgeTag(X.data,Y.data)
	               })
	    });
  
       assume(G.vertices.contains(root))
       root.data.values.foreach(r => {
         root.tag.asInstanceOf[VertexTag[A]].P(r) = P(root.data,r,Map()).get // P(r|a) = P(r) for the root
         root.tag.asInstanceOf[VertexTag[A]].pi(r) = P(root.data,r,Map()).get // pi(r) = P(r) for the root ... since there are no parents and known-children.. there is no one to consult with 
                                })

       // for each child X of R ... send pi message
       root.getEdges.foreach(e => send_pi_msg(root,e.getOther(root)))
  }
  
  
  def update_tree(V:Variable[A],_v:A) {
    V(Some(_v))
    aA = aA + V
    
    val vVertex = G.vertex(V).get
    
    val vTag = vVertex.tag.asInstanceOf[VertexTag[A]]
    vTag.pi(_v) = 1
    vTag.lambda(_v) = 1
    vTag.P(_v) = 1
    
    V.values.filter(_ != _v).foreach(v => {
			    vTag.pi(v) = 0
			    vTag.lambda(v) = 0
			    vTag.P(v) = 0
   })
    
    val Z = vVertex.inEdges.head.getOther(vVertex)
    if(vVertex != root && !aA.contains(Z.data)){
      send_lambda_msg(vVertex,Z)
    }
    
    vVertex.getEdges.foreach(e => {
      val X = e.getOther(vVertex)
      if(!aA.contains(X.data)){
        send_pi_msg(vVertex,X)
      }})
  }
  
  
  def send_lambda_msg(Y:Vertex[Variable[A]],X:Vertex[Variable[A]]) {
   
        val edge = X.getEdges.filter(_.getOther(X) == Y).head

        X.data.values.foreach(x => {
          // calculate Lambda values for each x ...lambda that has to come up to x from y
          // lambda-Y(x)  = summation-of(P(y|x)*lambda(y))
            val vtx = X.tag.asInstanceOf[VertexTag[A]]
          edge.tag.asInstanceOf[EdgeTag[A]].lambda(x) = Y.data.values.foldRight(0f)((y,l) => l + P(Y.data,y,Map((X.data,x))).get*Y.tag.asInstanceOf[VertexTag[A]].lambda(y))
          vtx.lambda(x) = X.getEdges.foldRight[float](1)((e,lambda) => (lambda) * e.tag.asInstanceOf[EdgeTag[A]].lambda(x))
          
          vtx.P(x) = vtx.pi(x)*vtx.lambda(x) // this has a hidden alpha .. hence the need to normalize later

        })
        val p = X.tag.asInstanceOf[VertexTag[A]].P 
            //NORMALIZE P(x|a) = P(x1|a) / P(x1|a) +P(x2|a)...+P(xn|a) 
	    val total = p.values./:(0f)(_ + _)
	    p.keySet.foreach(x => {p(x) = p(x)/total})
        
     
       // ASSUMING that it's a tree structure ... X.inedges would have only one member and would also
     
     
     if(root != X && !aA.contains(X.inEdges.head.getOther(X).data))
       send_lambda_msg(X,X.inEdges.head.getOther(X))
     

     X.getEdges.filter(e => e.getOther(X)!= Y && !aA.contains(e.getOther(X).data)).foreach(
       W => send_pi_msg(X,W.getOther(X))
     )
    
    
  }
  
  
  def send_pi_msg(Z:Vertex[Variable[A]],X:Vertex[Variable[A]]){
    
    //get the edge that connects Z to X (there could be better ways of getting this edge... say directly from the graph ... 
    val edge = G.getEdge(Z,X).get
    
    
    Z.data.values.foreach(z => {//for each possible value z of Z
      // calculate the product of all the Lambda messages coming up from the children of Z (excluding X)
      val lambdaMsgs = Z.getEdges.filter(_.getOther(Z) != X).foldRight[float](1f)((e,lm) => lm*e.tag.asInstanceOf[EdgeTag[A]].lambda(z))
      //calculate the pi message for X for z as a product of pi value of z and the product of lambda messages calculated above
      edge.tag.asInstanceOf[EdgeTag[A]].pi(z) = Z.tag.asInstanceOf[VertexTag[A]].pi(z) * lambdaMsgs
      
      
	})

    
    
    X.data.values.foreach(x => {//for each value x of X
      
      val vt = X.tag.asInstanceOf[VertexTag[A]]
      //compute X's pi value for x
         //pi(x) is summation of the prduct of P(x|z) and piX(z) for all z
       vt.pi(x)  = Z.data.values.foldRight[float](0f)((z,singlePi) => singlePi + P(X.data,x,Map((Z.data,z))).get*edge.tag.asInstanceOf[EdgeTag[A]].pi(z))
       // compute X's probability for x|a
       
       vt.P(x) = vt.pi(x)*vt.lambda(x) // this has a hidden alpha .. hence the need to normalize later
       
       
       
	})
    
    val p = X.tag.asInstanceOf[VertexTag[A]].P 
    
    //NORMALIZE P(x|a) = P(x1|a) / P(x1|a) +P(x2|a)...+P(xn|a) 
    val total = p.values./:(0f)(_ + _)
    p.keySet.foreach(x => {p(x) = p(x)/total})
 
   
    //for each child Y of X such that Y is not an
     X.getEdges.filter(e => !aA.contains(e.getOther(X).data)).foreach(e =>{
       val Y = e.getOther(X)
           send_pi_msg(X,Y)
       }
     )
    
  }
  
  
}



